"""
多分类问题：mnist手写数字(全链接神经网络)
"""
from keras.utils import to_categorical
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense
from utils.plot_graph import keras_plot_model
from utils.search_model_paramter import grid_search
from utils.model_fit import train_model


def get_data():
    """

    :return:
    """
    (train_images, train_label), (test_images, test_label) = mnist.load_data()
    # (?, 28, 28)
    train_images = \
        train_images.reshape((train_images.shape[0],
                              train_images.shape[1] * train_images.shape[2]))
    train_images = train_images.astype(float)
    train_images /= 255
    # 将整型标签转为onehot
    train_label = to_categorical(train_label)

    test_images = \
        test_images.reshape((test_images.shape[0],
                             test_images.shape[1] * test_images.shape[2]))
    test_images = test_images.astype(float)
    test_images /= 255

    test_label = to_categorical(test_label)

    return train_images, train_label, test_images, test_label


def build_fit_model(train_x, train_y, test_x, test_y,
                    unit, act, opt, loss, epochs, batch_size, verbose,
                    pic_path, model_name, workers=4,
                    use_multiprocessing=True, shuffle=True, plot=False):
    """

    :param train_x:
    :param train_y:
    :param test_x:
    :param test_y:
    :param unit:
    :param act:
    :param opt:
    :param loss:
    :param epochs:
    :param batch_size:
    :param verbose:
    :param pic_path:
    :param model_name:
    :param workers:
    :param use_multiprocessing:
    :param shuffle:
    :param plot:
    :return:
    """
    model = Sequential()
    model.add(Dense(units=unit,
                    activation=act[0],
                    input_shape=(train_x.shape[1], )))
    model.add(Dense(units=train_y.shape[1], activation=act[1]))

    model.compile(optimizer=opt,
                  loss=loss,
                  metrics=['accuracy'])

    return train_model(model=model,
                       train_x=train_x, train_y=train_y,
                       test_x=test_x, test_y=test_y,
                       epochs=epochs,
                       model_name=model_name,
                       pic_path=pic_path,
                       shuffle=shuffle,
                       verbose=verbose,
                       plot=plot,
                       batch_size=batch_size,
                       workers=workers,
                       use_multiprocessing=use_multiprocessing)


def multi_class_full_nn_run(workers):
    """

    :return:
    """
    get_parameter = False
    train_model = True

    unit_arr = [64, 128, 256, 512, 1024]
    # act_arr = ['softmax', 'relu']
    act_arr = [['softmax', 'softmax'], ['softmax', 'relu'],
               ['relu', 'relu'], ['relu', 'softmax']]
    opt_arr = ['RMSprop', 'Adam', 'Adadelta', 'Nadam']
    loss_arr = ['categorical_crossentropy']
    epochs = 20
    batch_size = 300

    train_images, train_label, test_images, test_label = get_data()
    # 参数选择
    if get_parameter:
        grid_search(func=build_fit_model,
                    train_x=train_images, train_y=train_label,
                    test_x=test_images, test_y=test_label,
                    unit_arr=unit_arr,
                    act_arr=act_arr,
                    opt_arr=opt_arr,
                    loss_arr=loss_arr,
                    epochs=epochs,
                    batch_size=batch_size,
                    workers=workers,
                    use_multiprocessing=True,
                    shuffle=True,
                    plot=False)

    # 模型训练
    if train_model:
        model, loss_metrics = \
            build_fit_model(train_x=train_images, train_y=train_label,
                            test_x=test_images, test_y=test_label,
                            unit=512,
                            act=['relu', 'softmax'],
                            # opt='RMSprop',
                            opt='Nadam',
                            loss='categorical_crossentropy',
                            pic_path='../pic',
                            model_name='多分类全链接神经网络',
                            epochs=epochs,
                            batch_size=batch_size,
                            verbose=1,
                            workers=workers,
                            use_multiprocessing=True,
                            shuffle=True,
                            plot=True)
        print('loss_metrics: {}'.format(loss_metrics))
        # plot model
        keras_plot_model(model=model, pic='../pic/多分类全链接神经网络.png')
